
import tushare as ts
import pandas as pd
import time
import datetime

#用于首次获取2年股票数据
def main():
    # 设置tushare pro的token并获取连接
    ts.set_token('883c1e2da92c0f68194fe2c4b4aa0e4da880b41ef62ff556f2da00e7')
    pro = ts.pro_api()

    #得到2年所有交易日
    datelist = pro.trade_cal(exchange='', start_date='20170101', end_date='20190124')
    datelist = datelist[datelist.is_open > 0]['cal_date']

    stocks_data = pd.DataFrame()

    i = 0
    j = 0#计数用

    #for循环遍历全部交易日
    for date in datelist:
        if i > 200:#tushare访问次数限制
            print('一分钟请求超过200，休息一分钟继续')
            time.sleep(60)
            i = 0
        else:
            print(date)
            df_data = pro.daily(trade_date=date)#获取每天个股数据
            df_factor = pro.adj_factor(ts_code='', trade_date=date)  # 获取每天个股复权因子

            #factor表 数据清洗
            #循环遍历因子表中全部股票代码，若其不在股票数据表中，删除之
            for code in df_factor['ts_code'].values.tolist():
                if code not in df_data['ts_code'].values.tolist():
                    print('删除factor表中多余的不在data表的元素')#如B股退市股
                    df_factor = df_factor.set_index(df_factor['ts_code'], drop=True)
                    df_factor.drop(code,axis = 0 ,inplace = True)

            #若factor表含有股票数据中所没有的股票，添加之
            if len(df_data) != len(df_factor ):
                print('factor表元素缺失，正在补齐...')
                j = j+1
                for code in  df_data['ts_code'].values.tolist():
                    if code not in df_factor['ts_code'].values.tolist():
                        df1 = pro.adj_factor(ts_code=code, trade_date=date)#获取缺失的股票当天的因子数据
                        df1 = df1.set_index(df1['ts_code'], drop=True)#按代码排序
                        df_factor = df_factor.append(df1)
           # df_factor.drop(['ts_code'], axis=1, inplace=True)
           # df_factor.reset_index(inplace=True)
            df_factor.drop(['ts_code', 'trade_date'], axis=1, inplace=True)
            df_data = df_data.set_index(df_data['ts_code'],drop=True).sort_index(level=0)
            #df_data.drop(['ts_code'], axis=1, inplace=True)
            df_data = df_data.join(df_factor)

            stocks_data = stocks_data.append(df_data)  # 把因子表赋值到股票数据中

        i = i+1

    # 创建多层索引"trade_date|ts_code"
    stocks_data = stocks_data.set_index([stocks_data['trade_date'],stocks_data['ts_code']], drop=True)
    stocks_data = stocks_data.drop('trade_date', 1)

    # 保存
    h5 = pd.HDFStore('stock_data.h5', 'w', complevel=4, complib='blosc')
    h5['data'] = stocks_data
    h5.close()
    print(j)


if __name__ == '__main__':
    main()

